package lcrypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"errors"
	"fmt"
	"strings"

	"github.com/spf13/viper"
	"github.com/zenazn/pkcs7pad"
	"go.uber.org/zap"
)

const (
	PKCS5 = "PKCS5"
	PKCS7 = "PKCS7"
)

func NewAES(logger *zap.Logger, conf *viper.Viper) *AES {
	model := strings.ToLower(strings.TrimSpace(conf.GetString("aes.model")))
	pad := strings.ToUpper(strings.TrimSpace(conf.GetString("aes.model")))
	key := conf.GetString("aes.key")
	iv := conf.GetString("aes.iv")

	switch model {
	case "cbc", "":
		model = "cbc"
		// AES分组长度为128位，所以blockSize=16，单位字节
		if len(iv) < 16 {
			logger.Fatal(
				"[初始化]aes加密初始化向量长度不足",
				zap.String("model", model),
				zap.String("iv", iv),
			)
		}
		iv = iv[:16]
		if pad != PKCS5 && pad != PKCS7 {
			logger.Fatal(
				"[初始化]aes加密补齐方式不正确",
				zap.String("model", model),
				zap.String("padding", pad),
			)
		}
	case "ecb":
		if pad != PKCS7 {
			logger.Fatal(
				"[初始化]aes加密补齐方式不正确",
				zap.String("model", model),
				zap.String("padding", pad),
			)
		}
	default:
		logger.Fatal("[初始化]不支持的aes加密模式", zap.String("模式", model))
	}

	return &AES{
		model: model,
		pad:   pad,
		key:   []byte(key),
		iv:    []byte(iv),
	}
}

type AES struct {
	model string
	pad   string
	key   []byte
	iv    []byte
}

func (a *AES) Name() string {
	return fmt.Sprintf("aes_%s", a.model)
}

func (a *AES) Encrypt(data []byte) ([]byte, error) {
	switch a.model {
	case "cbc":
		return a.CbcEncrypt(data)
	case "ecb":
		return a.EcbEncrypt(data)
	default:
		return nil, errors.New("不支持的模式")
	}
}

func (a *AES) Decrypt(data []byte) ([]byte, error) {
	switch a.model {
	case "cbc":
		return a.CbcDecrypt(data)
	case "ecb":
		return a.EcbDecrypt(data)
	default:
		return nil, errors.New("不支持的模式")
	}
}

// AESCBCEncrypt AES CBC加密
func (a *AES) CbcEncrypt(plain []byte) ([]byte, error) {
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}

	var paddedData []byte

	// AES分组长度为128位，所以blockSize=16，单位字节
	blockSize := block.BlockSize()

	if len(a.iv) > 0 && len(a.iv) != blockSize {
		return nil, errors.New("iv length error")
	}

	switch a.pad {
	case PKCS5:
		paddedData = PKCS5Pad(plain, blockSize)
	case PKCS7:
		paddedData = pkcs7pad.Pad(plain, blockSize)
	default:
		return nil, errors.New("padding error")
	}
	blockMode := cipher.NewCBCEncrypter(block, a.iv) // 初始向量的长度必须等于块block的长度16字节
	cipher := make([]byte, len(paddedData))
	blockMode.CryptBlocks(cipher, paddedData)
	return cipher, nil
}

// AESCBCDecrypt AES CBC解密
func (a *AES) CbcDecrypt(crypted []byte) (plain []byte, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("%v", e)
		}
	}()
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return
	}

	if len(a.iv) != block.BlockSize() {
		return nil, errors.New("iv length error")
	}

	blockMode := cipher.NewCBCDecrypter(block, a.iv) // 初始向量的长度必须等于块block的长度16字节
	paddedData := make([]byte, len(crypted))
	blockMode.CryptBlocks(paddedData, crypted)

	switch a.pad {
	case PKCS5:
		return PKCS5Unpad(paddedData), nil
	case PKCS7:
		return pkcs7pad.Unpad(paddedData)
	default:
		return nil, errors.New("padding error")
	}
}

// AESEcbEncrypt ECB加密
func (a *AES) EcbEncrypt(plain []byte) ([]byte, error) {
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}

	blockSize := block.BlockSize()
	plain = pkcs7pad.Pad(plain, blockSize)
	cipher := make([]byte, len(plain))

	for bs, be := 0, blockSize; bs < len(plain); bs, be = bs+be, be+be {
		block.Encrypt(cipher[bs:be], plain[bs:be])
	}

	return cipher, nil
}

// AESEcbDecrypt ECB解密
func (a *AES) EcbDecrypt(cipher []byte) ([]byte, error) {
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}

	plain := make([]byte, len(cipher))
	for bs, be := 0, block.BlockSize(); bs < len(plain); bs, be = bs+be, be+be {
		block.Decrypt(plain[bs:be], cipher[bs:be])
	}

	return pkcs7pad.Unpad(plain)
}

// PKCS5Padding 填充明文
func PKCS5Pad(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// PKCS5Unpad 去除填充数据
func PKCS5Unpad(data []byte) []byte {
	length := len(data)
	unPadding := int(data[length-1])
	return data[:(length - unPadding)]
}
